{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入优化后的工具模块\n",
    "from utils import get_latest_ssq_data, save_to_csv\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "import logging\n",
    "\n",
    "# 配置日志\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# 设置显示选项\n",
    "pd.set_option('display.float_format', '{:.2f}'.format)\n",
    "plt.style.use('seaborn')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据获取与预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_prepare_data():\n",
    "    \"\"\"\n",
    "    加载并预处理双色球数据\n",
    "    \n",
    "    Returns:\n",
    "        pd.DataFrame: 处理后的数据\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 使用优化后的工具函数获取数据\n",
    "        df = get_latest_ssq_data(days=365*3)  # 获取3年数据\n",
    "        \n",
    "        if df.empty:\n",
    "            raise ValueError(\"获取的数据为空\")\n",
    "            \n",
    "        # 数据排序\n",
    "        df = df.sort_values('code', ascending=True)\n",
    "        \n",
    "        # 保存数据\n",
    "        save_to_csv(df, 'lottery_data.csv')\n",
    "        \n",
    "        return df\n",
    "        \n",
    "    except Exception as e:\n",
    "        logger.error(f\"数据加载失败: {e}\")\n",
    "        return pd.DataFrame()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据分析与可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyze_number_frequency(df):\n",
    "    \"\"\"\n",
    "    分析号码出现频率并可视化\n",
    "    \n",
    "    Args:\n",
    "        df: 包含双色球数据的DataFrame\n",
    "    \"\"\"\n",
    "    if df.empty:\n",
    "        return\n",
    "        \n",
    "    try:\n",
    "        # 红球频率分析\n",
    "        red_cols = ['red1', 'red2', 'red3', 'red4', 'red5', 'red6']\n",
    "        red_counts = df[red_cols].stack().value_counts().sort_index()\n",
    "        \n",
    "        # 蓝球频率分析\n",
    "        blue_counts = df['blue'].value_counts().sort_index()\n",
    "        \n",
    "        # 可视化\n",
    "        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))\n",
    "        \n",
    "        # 红球分布\n",
    "        sns.barplot(x=red_counts.index, y=red_counts.values, ax=ax1, palette='Reds_r')\n",
    "        ax1.set_title('红球号码出现频率')\n",
    "        ax1.set_xlabel('号码')\n",
    "        ax1.set_ylabel('出现次数')\n",
    "        \n",
    "        # 蓝球分布\n",
    "        sns.barplot(x=blue_counts.index, y=blue_counts.values, ax=ax2, palette='Blues_r')\n",
    "        ax2.set_title('蓝球号码出现频率')\n",
    "        ax2.set_xlabel('号码')\n",
    "        ax2.set_ylabel('出现次数')\n",
    "        \n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "        \n",
    "        return red_counts, blue_counts\n",
    "        \n",
    "    except Exception as e:\n",
    "        logger.error(f\"频率分析失败: {e}\")\n",
    "        return None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 基于频率的预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_by_frequency(red_counts, blue_counts):\n",
    "    \"\"\"\n",
    "    基于频率预测下一期号码\n",
    "    \n",
    "    Args:\n",
    "        red_counts: 红球频率数据\n",
    "        blue_counts: 蓝球频率数据\n",
    "    \"\"\"\n",
    "    if red_counts is None or blue_counts is None:\n",
    "        return\n",
    "        \n",
    "    try:\n",
    "        # 转换为PyTorch张量\n",
    "        red_probs = torch.tensor(red_counts.values / red_counts.sum())\n",
    "        blue_probs = torch.tensor(blue_counts.values / blue_counts.sum())\n",
    "        \n",
    "        # 预测最高频号码\n",
    "        _, top_red_indices = red_probs.sort(descending=True)\n",
    "        top_red = (top_red_indices[:6] + 1).tolist()\n",
    "        top_blue = (blue_probs.argmax() + 1).item()\n",
    "        \n",
    "        # 预测最低频号码\n",
    "        low_red = (top_red_indices[-6:] + 1).tolist()\n",
    "        low_blue = (blue_probs.argmin() + 1).item()\n",
    "        \n",
    "        print(f\"预测的6个红球号码(高频): {sorted(top_red)}\")\n",
    "        print(f\"预测的蓝球号码(高频): {top_blue}\")\n",
    "        print(f\"预测的6个红球号码(低频): {sorted(low_red)}\")\n",
    "        print(f\"预测的蓝球号码(低频): {low_blue}\")\n",
    "        \n",
    "    except Exception as e:\n",
    "        logger.error(f\"频率预测失败: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 基于LSTM的预测模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LotteryDataset(Dataset):\n",
    "    \"\"\"双色球数据集类\"\"\"\n",
    "    def __init__(self, X, y):\n",
    "        self.X = X\n",
    "        self.y = y\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        return self.X[idx], self.y[idx]\n",
    "\n",
    "class LotteryPredictor(nn.Module):\n",
    "    \"\"\"双色球预测模型\"\"\"\n",
    "    def __init__(self, input_size, hidden_size, num_layers, output_size, dropout=0.2):\n",
    "        super().__init__()\n",
    "        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout)\n",
    "        self.fc = nn.Linear(hidden_size, output_size)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        out, _ = self.lstm(x)\n",
    "        out = self.fc(out[:, -1, :])\n",
    "        return out\n",
    "\n",
    "def train_lstm_model(df, target='red', epochs=100):\n",
    "    \"\"\"\n",
    "    训练LSTM预测模型\n",
    "    \n",
    "    Args:\n",
    "        df: 包含双色球数据的DataFrame\n",
    "        target: 预测目标('red'或'blue')\n",
    "        epochs: 训练轮数\n",
    "    \"\"\"\n",
    "    if df.empty:\n",
    "        return None, None\n",
    "        \n",
    "    try:\n",
    "        # 数据准备\n",
    "        if target == 'red':\n",
    "            data = df[['red1', 'red2', 'red3', 'red4', 'red5', 'red6']].values\n",
    "            output_size = 6\n",
    "        else:\n",
    "            data = df[['blue']].values\n",
    "            output_size = 1\n",
    "            \n",
    "        # 归一化\n",
    "        scaler = MinMaxScaler()\n",
    "        data_normalized = scaler.fit_transform(data)\n",
    "        \n",
    "        # 创建时间序列数据\n",
    "        def create_sequences(data, window_size=10):\n",
    "            X, y = [], []\n",
    "            for i in range(len(data)-window_size):\n",
    "                X.append(data[i:i+window_size])\n",
    "                y.append(data[i+window_size])\n",
    "            return np.array(X), np.array(y)\n",
    "            \n",
    "        X, y = create_sequences(data_normalized)\n",
    "        \n",
    "        # 划分训练测试集\n",
    "        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n",
    "        \n",
    "        # 转换为张量\n",
    "        X_train = torch.FloatTensor(X_train)\n",
    "        X_test = torch.FloatTensor(X_test)\n",
    "        y_train = torch.FloatTensor(y_train)\n",
    "        y_test = torch.FloatTensor(y_test)\n",
    "        \n",
    "        # 创建数据加载器\n",
    "        train_dataset = LotteryDataset(X_train, y_train)\n",
    "        test_dataset = LotteryDataset(X_test, y_test)\n",
    "        \n",
    "        train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "        test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)\n",
    "        \n",
    "        # 初始化模型\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        model = LotteryPredictor(\n",
    "            input_size=X_train.shape[2], \n",
    "            hidden_size=64, \n",
    "            num_layers=2, \n",
    "            output_size=output_size\n",
    "        ).to(device)\n",
    "        \n",
    "        criterion = nn.MSELoss()\n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "        \n",
    "        # 训练循环\n",
    "        best_loss = float('inf')\n",
    "        for epoch in range(epochs):\n",
    "            model.train()\n",
    "            for inputs, targets in train_loader:\n",
    "                inputs, targets = inputs.to(device), targets.to(device)\n",
    "                \n",
    "                optimizer.zero_grad()\n",
    "                outputs = model(inputs)\n",
    "                loss = criterion(outputs, targets)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                \n",
    "            # 验证\n",
    "            model.eval()\n",
    "            with torch.no_grad():\n",
    "                val_loss = 0\n",
    "                for inputs, targets in test_loader:\n",
    "                    inputs, targets = inputs.to(device), targets.to(device)\n",
    "                    outputs = model(inputs)\n",
    "                    val_loss += criterion(outputs, targets).item()\n",
    "                \n",
    "                val_loss /= len(test_loader)\n",
    "                \n",
    "                if val_loss < best_loss:\n",
    "                    best_loss = val_loss\n",
    "                    best_model = model.state_dict()\n",
    "                    \n",
    "            if (epoch+1) % 10 == 0:\n",
    "                print(f'Epoch [{epoch+1}/{epochs}], Loss: {val_loss:.4f}')\n",
    "                \n",
    "        print(f'Training completed. Best validation loss: {best_loss:.4f}')\n",
    "        model.load_state_dict(best_model)\n",
    "        \n",
    "        return model, scaler\n",
    "        \n",
    "    except Exception as e:\n",
    "        logger.error(f\"模型训练失败: {e}\")\n",
    "        return None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 主流程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    \"\"\"主执行流程\"\"\"\n",
    "    # 1. 数据加载\n",
    "    df = load_and_prepare_data()\n",
    "    if df.empty:\n",
    "        return\n",
    "    \n",
    "    # 2. 数据分析\n",
    "    red_counts, blue_counts = analyze_number_frequency(df)\n",
    "    \n",
    "    # 3. 基于频率的预测\n",
    "    predict_by_frequency(red_counts, blue_counts)\n",
    "    \n",
    "    # 4. 训练红球预测模型\n",
    "    red_model, red_scaler = train_lstm_model(df, target='red')\n",
    "    \n",
    "    # 5. 训练蓝球预测模型\n",
    "    blue_model, blue_scaler = train_lstm_model(df, target='blue')\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
